# -*- coding:utf-8 -*-
from engine.data.dataloader import D
from sklearn.model_selection import train_test_split
from sklearn.metrics import r2_score
import lightgbm as lgb
from mlflow import log_metric, log_param, log_artifacts

class Dataset:
    def __init__(self, codes, handler):
        self.handler = handler
        fields, names = self.handler.get_feature_config()
        label_expr, label_name = self.handler.get_label_config()
        self.load_internal(codes, fields, names, label_expr, label_name)


    def load_internal(self, codes, fields, names, label_expr, label_name):
        if not names:
            names = fields
        all_fields = fields.copy()
        all_fields.append(label_expr)

        all_names = names.copy()
        all_names.append(label_name)

        self.df = D.load(codes, fields=all_fields, names=all_names)
        self.feature_names = names
        self.label_name = label_name


        #self.df[label_name] = self.df[label_name].groupby("date").apply(lambda x: (x - x.mean()).div(x.std()))



    #def __init__(self, codes, fields: list, label_expr, names=None,label_name='label'):
    #    self.load_internal(codes,fields, names, label_expr, label_name)


    def split(self, feature_names=None, test_size=0.3, shuffle=True):
        if not feature_names:
            feature_names = self.feature_names

        df_feature = self.df[feature_names]
        df_label = self.df[self.label_name]

        X_train, X_valid, y_train, y_valid = train_test_split(df_feature, df_label, test_size=test_size, shuffle=shuffle)
        return X_train, X_valid, y_train, y_valid


    def get_data(self, date_range=['20080101','20141231']):
        sub_df = self.df.loc[date_range[0]:date_range[1]]
        return sub_df[self.feature_names], sub_df['label']



if __name__ == '__main__':
    from engine.data.datahandler import DataHandler
    fields = ['Return($close,5)', 'Return($close,20)', 'Ref($close,126)/$close -1']
    names = ['return_5', 'return_20', 'return_126']

    label_expr = 'Ref($close,-5)/$close -1'
    label_expr = 'QCut(Ref($close,-5)/$close -1,5)'

    codes = ['000300.SH', '000905.SH']
    codes = [
        '000300.SH',

    ]
    #ds = Dataset(codes=codes, fields=fields, names=names,label_expr=label_expr)

    ds = Dataset(codes=codes, handler=DataHandler())
    print(ds.df[['label','label_10','date','code']])
    #print(ds.split())



